每天的專案會同步到 GitLab 上,可以前往 GitLab 查看,有興趣的朋友歡迎留言 or 來信討論,我的信箱是 nickchen1998@gmail.com。
ParentDocumentRetriever 的用途是平衡文件拆分和檢索的需求,適用於當文件過大時,無法直接將整個文件進行檢索的情境。它的主要功能是先將文件拆分成較小的片段來建立向量索引,但在檢索時,會回傳這些小片段的父文件(即原本較大的文件或一部分)。
與昨天的 metadata 相反,ParentDocumentRetriever 主要是想要解決文本內文不足的問題,由於過長的 chunk 可能會導致搜尋精度降低,因此時常需要將文件拆分成較小的片段來進行檢索。但這樣做可能會使得檢索結果過於分散,無法有效地回答使用者的問題,即便有使用 Recursion Strategy 也仍舊只能保留有限的文本,然而隨著語言模型能支援的 token 長度越來越大,也相對的可以給予更多的內容讓語言模型進行摘要回答。
from env_settings import EnvSettings, BASE_DIR
from langchain.retrievers import ParentDocumentRetriever
from langchain_openai.embeddings import OpenAIEmbeddings
from langchain_pinecone import PineconeVectorStore
from langchain_community.document_loaders import PyMuPDFLoader
from pinecone import Pinecone
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.storage import InMemoryStore
# 準備兩個 Pinecone 連線
env_settings = EnvSettings()
vector_store = PineconeVectorStore(
index=Pinecone(api_key=env_settings.PINECONE_API_KEY).Index("chunks"),
embedding=OpenAIEmbeddings(openai_api_key=env_settings.OPENAI_API_KEY)
)
# 準備文本切割器
bic_chunk_spliter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=50)
small_chunk_spliter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=50)
# 讀取文件並建立檢索器
loader = PyMuPDFLoader(file_path=BASE_DIR / "勞動基準法.pdf")
retriever = ParentDocumentRetriever(
vectorstore=vector_store,
docstore=InMemoryStore(),
child_splitter=small_chunk_spliter,
parent_splitter=bic_chunk_spliter,
)
# 寫入文件
retriever.add_documents(loader.load())
在上面的程式碼當中我們可以看到在準備切割器的時候準備了兩組,分別是 1000 字以及 200 字的,
並且我們呼叫了了 ParentDocumentRetriever
這個類別,並且將 vectorstore
與 docstore
進行初始化,這邊要特別注意的是,我們先使用了 InMemoryStore
這個類別作為文本的暫存,稍後會為大家說明該如何套用到自家的資料庫上。
在同一段程式碼的結尾我們加上下方這三行程式碼,這樣我們就可以透過問題來檢索段落了。
# 提問並檢索段落
question = "我可以因為不喜歡某個勞工而終止勞動契約嗎?"
retrieved_docs = retriever.invoke(question)
print(retrieved_docs[0].page_content)
下圖為搜尋出來的結果,可以看到我們成功的透過跟昨天一樣的問題,將更長的相關文本檢索出來:
剛剛在進行 add_documents()
動作的時候,應該可以發現我們有成功的把向量資料寫入到 Pinecone 中,但是我們的文本資料並沒有被寫入到任何地方,這也是為什麼檢索這段程式是請大家接著寫,而不是另外寫一段,這邊我們可以透過 docstore
來將文本資料持久化到我們的資料庫中,這樣就可以讓我們在下次檢索的時候不要重新讀取文件,而是直接從資料庫中讀取。
from langchain.storage.base import BaseStore
from pymongo import MongoClient
from typing import List, Optional, Sequence, Tuple, Union, Iterator
class MongoDBStore(BaseStore[str, str]):
def __init__(self, db_name="document_db", collection_name="documents"):
# 初始化 MongoDB 客戶端和資料庫
self.client = MongoClient("mongodb://localhost:27017/")
self.db = self.client[db_name]
self.collection = self.db[collection_name]
def mget(self, keys: Sequence[str]) -> List[Optional[str]]:
"""根據鍵值批量檢索文件內容"""
documents = self.collection.find({"_id": {"$in": keys}})
result_dict = {doc["_id"]: doc["document"] for doc in documents}
return [result_dict.get(key) for key in keys]
def mset(self, key_value_pairs: Sequence[Tuple[str, str]]) -> None:
"""批量添加或更新文件"""
for key, value in key_value_pairs:
self.collection.update_one(
{"_id": key},
{"$set": {"document": value}},
upsert=True
)
def mdelete(self, keys: Sequence[str]) -> None:
"""批量刪除文件"""
self.collection.delete_many({"_id": {"$in": keys}})
def yield_keys(self, prefix: Optional[str] = None) -> Iterator[str]:
"""返回鍵值的迭代器"""
if prefix:
cursor = self.collection.find({"_id": {"$regex": f"^{prefix}"}})
else:
cursor = self.collection.find()
for doc in cursor:
yield doc["_id"]
這邊我們實作了一個簡單的 MongoDBStore 類別,這個類別繼承了 BaseStore 並實作了 mget、mset、mdelete 以及 yield_keys 四個方法,這四個方法分別是用來批量檢索、批量添加或更新、批量刪除以及返回鍵值的迭代器,這樣我們就可以將文本資料持久化到 MongoDB 中了。
如此一來我們就可以將整體的流程拆分成,「計算」、「儲存」、「檢索」三個步驟,讓我們的程式碼更加模組化,也更容易維護。
今天我們介紹了 ParentDocumentRetriever 這個檢索器,透過這個檢索器我們可以將文件拆分成較小的片段來進行大段落的檢索,可以補足向量文本過短的問題,明天我們要來介紹 MultiQueryRetriever 這個檢索器。